package extraction;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import com.hp.hpl.jena.query.QueryExecution;
import com.hp.hpl.jena.query.QueryExecutionFactory;
import com.hp.hpl.jena.query.QuerySolution;
import com.hp.hpl.jena.query.ResultSet;
import com.hp.hpl.jena.rdf.model.InfModel;
import com.hp.hpl.jena.rdf.model.Model;
import com.hp.hpl.jena.rdf.model.ModelFactory;
import com.hp.hpl.jena.rdf.model.Property;
import com.hp.hpl.jena.rdf.model.RDFNode;
import com.hp.hpl.jena.rdf.model.Resource;
import com.hp.hpl.jena.rdf.model.Statement;
import com.hp.hpl.jena.rdf.model.StmtIterator;
import com.hp.hpl.jena.reasoner.rulesys.GenericRuleReasoner;
import com.hp.hpl.jena.reasoner.rulesys.Rule;
import com.hp.hpl.jena.tdb.TDBFactory;
import com.hp.hpl.jena.tdb.base.file.Location;

public class Inference {
	
	static Model model;
	static String prefix = "http://dbpedia.org/ontology/";
	static Set<String> relations;
		
	public static void inference(String TDBlocation,String rulesFile) throws IOException {
		
		/* load a stored Model */
		String directory = TDBlocation;
		System.out.println("Loading TBD from " + TDBlocation);
		Location loc = new Location(directory);		
		model = TDBFactory.createModel(loc);
		relations = new HashSet<String>();
		query();
		
		/* write original relations to file */
		System.out.println("Dumping original relations ... ");
		for (String r : relations) listAll(prefix, r, model, false);
		
        /* infer new relations */
		Model dedu = inference(model,rulesFile);
		
		/* write inferred relations to file */
		System.out.println("Dumping infered relations ... ");
		for (String r : relations) listAll(prefix, r, dedu, true);
		
		model.close();
	}
	
    public static void query() {
    	String query =  "SELECT DISTINCT ?p WHERE { ?s ?p ?o }";    	
    	QueryExecution qExe = QueryExecutionFactory.create(query, model);    	
    	ResultSet results = qExe.execSelect();
    	RDFNode rdf = null;
    	String relation;
    	while (results.hasNext()) {
    		QuerySolution r = results.next();    		
    		rdf = r.get("p"); 
    		relation = rdf.toString();
    		relations.add(relation.replaceFirst(prefix, ""));
    	}    	    	
    }
    
	public static void listAll(String prefix, String relation, Model model, boolean inferred) throws IOException{
		StmtIterator stmItr = model.listStatements();
		String filename;		
		
		if (!inferred) filename = relation.replaceFirst(prefix, "")+"_original.n3";
		else filename = relation.replaceFirst(prefix, "")+"_inferred.n3";
		
		BufferedWriter out = new BufferedWriter(new FileWriter(new File(filename)));
		
		while (stmItr.hasNext()) {			
			Statement stmt      = stmItr.nextStatement();  	// get next statement
		    Resource  subject   = stmt.getSubject();     	// get the subject
		    Property  predicate = stmt.getPredicate();   	// get the predicate
		    RDFNode   object    = stmt.getObject();      	// get the object
		    
		    if (predicate.toString().equals(prefix+relation)) {
		    	out.write(subject.toString());
		    	out.write(" " + predicate.toString() + " ");
			    if (object instanceof Resource) {
			       out.write(object.toString());
			    } else {
			        // object is a literal
			    	out.write(" \"" + object.toString() + "\"");
			    }
			out.write(" .\n"); 
			}
		}
		out.close();
	}
    
	public static Model inference(Model model,String rulesFile) {
		
		/* load rules */
        List<Rule> rules = Rule.rulesFromURL(rulesFile);
        
        System.out.println("Starting inference with the following rules: \n");
        for (Rule rule : rules) {
			System.out.println(rule.toString());
		}
        
        System.out.println();
        
        /* load reasoner */
        GenericRuleReasoner reasoner = new GenericRuleReasoner(rules);
        reasoner.setMode(GenericRuleReasoner.HYBRID);
        
        // load an Inference Model, from reasoner and ontology        
        InfModel inf = ModelFactory.createInfModel(reasoner, model);
        
        System.out.println("Applying rules ...");
        
        // make the inference!
        Model dedu = inf.getDeductionsModel();
        
        return dedu;
	}
	
}
